﻿// Copyright (c) 2021 raoyutian Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License. 
using Emgu.CV;
using System.Drawing;
using System.Text;
namespace PaddleOCR.Onnx.Paddle
{
    internal class OCREngine : IDisposable
    {
        #region 属性 
        internal OCRParameter Parameter { get; private  set; }
        private DBDetector det { get; set; }
        private Classifier cls { get; set; }
        private CRNNRecognizer rec { get; set; }
        #endregion

        public OCREngine(string det_infer, string cls_infer, string rec_infer, string keys, OCRParameter parameter)
        {
            det = new DBDetector();
            cls = new Classifier();
            rec = new CRNNRecognizer();
           if(parameter==null) parameter=new OCRParameter();
            Parameter = parameter;
            det.InitModel(det_infer, parameter.numThread);
            cls.InitModel(cls_infer, parameter.numThread);
            rec.InitModel(rec_infer, keys, parameter.numThread);
        }
        /// <summary>
        ///文本识别
        /// </summary>
        /// <param name="mat">Mat图像</param>
        /// <param name="parameter">参数</param>
        /// <returns></returns>
        public OCRResult Detect(Mat mat)
        {
            if (mat == null) throw new ArgumentNullException("mat");
            
            Mat originSrc = mat;
            int originMaxSide = Math.Max(originSrc.Cols, originSrc.Rows);

            int resize;
            if (Parameter.MaxSideLen <= 0 || Parameter.MaxSideLen > originMaxSide)
            {
                resize = originMaxSide;
            }
            else
            {
                resize = Parameter.MaxSideLen;
            }
            resize += 2 * Parameter.Padding;
            Rectangle paddingRect = new Rectangle(Parameter.Padding, Parameter.Padding, originSrc.Cols, originSrc.Rows);
            Mat paddingSrc = Utils.MakePadding(originSrc, Parameter.Padding);
            ScaleParam scale = ScaleParam.GetScaleParam(paddingSrc, resize);
            return DetectOnce(paddingSrc, paddingRect, scale, Parameter);
        }

        private OCRResult DetectOnce(Mat src, Rectangle originRect, ScaleParam scale, OCRParameter parameter)
        {
            if (parameter == null) parameter = new OCRParameter();

            Mat textBoxPaddingImg = src.Clone();
            int thickness = Utils.GetThickness(src);
           
            var startTicks = DateTime.Now.Ticks;

         
            var textBoxes = det.GetTextBoxes(src, scale, parameter.BoxScoreThresh, parameter.BoxThresh, parameter.UnClipRatio);
            var dbNetTime = (DateTime.Now.Ticks - startTicks) / 10000F;

            Utils.DrawTextBoxes(textBoxPaddingImg, textBoxes, thickness);
            List<Mat> partImages = Utils.GetPartImages(src, textBoxes);

            List<Angle> angles = cls.GetAngles(partImages, parameter.DoAngle, parameter.MostAngle);

            for (int i = 0; i < partImages.Count; ++i)
            {
                if (angles[i].Index == 1)
                {
                    partImages[i] = Utils.MatRotateClockWise180(partImages[i]);
                }
                
            }

            List<TextLine> textLines = rec.GetTextLines(partImages);
            
            List<TextBlock> textBlocks = new List<TextBlock>();
            for (int i = 0; i < textLines.Count; ++i)
            {
                TextBlock textBlock = new TextBlock();
                textBlock.BoxPoints = textBoxes[i].Points;
                textBlock.BoxScore = textBoxes[i].Score;
                textBlock.AngleIndex = angles[i].Index;
                textBlock.AngleScore = angles[i].Score;
                textBlock.AngleTime = angles[i].Time;
                textBlock.Text = textLines[i].Text;
                textBlock.CharScores = textLines[i].CharScores;
                textBlock.CrnnTime = textLines[i].Time;
                textBlock.BlockTime = angles[i].Time + textLines[i].Time;
                textBlocks.Add(textBlock);
            }
         
            var endTicks = DateTime.Now.Ticks;
            var fullDetectTime = (endTicks - startTicks) / 10000F;
           
           
            Mat boxImg = new Mat(textBoxPaddingImg, originRect);

            StringBuilder strRes = new StringBuilder();
            textBlocks.ForEach(x => strRes.AppendLine(x.Text));

            OCRResult ocrResult = new OCRResult();
            ocrResult.TextBlocks = textBlocks;
            ocrResult.DbNetTime = dbNetTime;
            ocrResult.BoxImg = boxImg;
            ocrResult.DetectTime = fullDetectTime;
            ocrResult.Text = string.Join("",textBlocks.Select(x=>x.Text).ToArray());

            return ocrResult;
        }
        public void Dispose()
        {
            det = null;
            cls = null;
            rec = null;
        }
    }
}
